class Solution
{
public:
    vector<int> countVisitedNodes(vector<int> &edges)
    {
        int n = edges.size();
        // 每个节点有且仅有一个出度, 所以图的形式是若干子图, 每个子图有一个环, 和一些指向环的链

        vector<int> ufs(n);
        iota(ufs.begin(), ufs.end(), 0);
        function<int(int)> find = [&](int x)
        {
            return x == ufs[x] ? x : ufs[x] = find(ufs[x]);
        };
        function<void(int, int)> merge = [&](int x, int y)
        {
            ufs[find(x)] = find(y);
        };
        for (int i = 0; i < n; ++i)
        {
            merge(i, edges[i]);
        }

        unordered_map<int, unordered_set<int>> cycle; // 统计每个子图的环
        for (int node : edges)
        {
            if (cycle.count(find(node)))
            {
                continue;
            }
            int start = node;
            int cur = node;
            while (!cycle[find(node)].count(cur))
            {
                cycle[find(node)].insert(cur);
                cur = edges[cur];
            }
            while (start != cur)
            {
                cycle[find(node)].erase(start);
                start = edges[start];
            }
        }

        vector<int> distToCycle(n, -1); // 统计每个节点到环的距离
        for (auto &[_, nodes] : cycle)
        {
            for (int node : nodes)
            {
                distToCycle[node] = 0;
            }
        }

        function<int(int)> dfs = [&](int node)
        {
            if (distToCycle[node] != -1)
            {
                return distToCycle[node];
            }
            distToCycle[node] = dfs(edges[node]) + 1;
            return distToCycle[node];
        };

        vector<int> result(n);
        for (int i = 0; i < n; ++i)
        {
            if (distToCycle[i] == -1)
            {
                dfs(i);
            }
            result[i] = distToCycle[i] + cycle[find(i)].size();
        }
        return result;
    }
};